from pyvis.network import Network
from neo4j_utils.connect import Neo4jConnection
import os


def visualize_graph():
    # 连接到Neo4j数据库
    neo4j = Neo4jConnection("bolt://127.0.0.1:7687", "neo4j", "abcd1234")

    # 初始化有向图网络
    net = Network(
        notebook=False,
        width="1000px",
        height="800px",
        bgcolor="#ffffff",
        font_color="black",
        directed=True  # 明确设置为有向图，确保关系方向正确显示
    )

    try:
        # 从Neo4j读取节点和关系数据
        with neo4j.driver.session() as session:
            # 查询所有节点（使用elementId避免deprecated警告）
            nodes = session.run("MATCH (n) RETURN elementId(n) as id, labels(n) as labels, properties(n) as props")
            for node in nodes:
                node_id = node["id"]
                labels = node["labels"]
                props = node["props"]

                # 根据节点标签设置不同颜色
                if "Country" in labels:
                    color = "#3498db"  # 蓝色 - 国家节点
                elif "Capital" in labels:
                    color = "#e74c3c"  # 红色 - 首都节点
                elif "Village" in labels:
                    color = "#2ecc71"  # 绿色 - 地球村节点
                else:
                    color = "#9b59b6"  # 紫色 - 其他节点

                # 设置节点显示名称和悬停提示
                node_name = props.get("name", f"Node {node_id}")
                net.add_node(
                    node_id,
                    label=node_name,
                    color=color,
                    title=str(props)  # 悬停时显示节点属性
                )

            # 查询所有关系
            relationships = session.run(
                "MATCH (a)-[r]->(b) RETURN elementId(a) as source, elementId(b) as target, type(r) as type"
            )
            for rel in relationships:
                net.add_edge(
                    rel["source"],
                    rel["target"],
                    label=rel["type"],  # 显示关系类型标签
                    arrows="to"  # 明确显示箭头方向
                )

        # 配置布局参数（提升图形可读性）
        net.barnes_hut(
            gravity=-800,      # 节点间引力（负值为排斥）
            spring_length=200, # 边的长度
            damping=0.4        # 动画阻尼（稳定性）
        )

        # 生成并保存可视化HTML文件
        output_file = "knowledge_graph.html"
        net.write_html(output_file)

        # 验证文件生成
        if os.path.exists(output_file):
            print(f"知识图谱已保存至: {os.path.abspath(output_file)}")
        else:
            raise Exception(f"可视化文件生成失败: {output_file}")

    except Exception as e:
        print(f"可视化过程出错: {str(e)}")
    finally:
        # 确保数据库连接关闭
        neo4j.close()